{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e8b27860",
   "metadata": {},
   "source": [
    "# 拆分数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76bac66f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "from onekey_algo import get_param_in_cwd\n",
    "\n",
    "bagsize = get_param_in_cwd('bagsize')\n",
    "model_root = os.path.join(get_param_in_cwd('model_root'), 'semi_models', get_param_in_cwd('sel_model')['semi'])\n",
    "features = pd.read_csv(os.path.join(model_root, 'viz', f'dlfeatures.csv'), header=None)\n",
    "features.columns = ['ID'] + [f'DL_{i}' for i in range(features.shape[-1] - 1)]\n",
    "all_features = features\n",
    "all_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f99cf8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ID(x):\n",
    "    items = x.split('_')\n",
    "    return '_'.join(items[:-6])\n",
    "all_features['ID'] = all_features['ID'].map(lambda x: get_ID(x))\n",
    "all_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f28566a",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_2048_features = []\n",
    "for sid in set(all_features['ID']):\n",
    "    sample_features = all_features[all_features['ID'] == sid]\n",
    "    if sample_features.shape[0] > bagsize:\n",
    "        sample_features = sample_features.sample(n=bagsize)\n",
    "    all_2048_features.append(sample_features)\n",
    "all_2048_features = pd.concat(all_2048_features, axis=0)\n",
    "all_2048_features.to_csv(os.path.join(model_root, f'dl{bagsize}.csv'), index=False)\n",
    "all_2048_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "758959b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from onekey_algo import get_param_in_cwd\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "label_data = pd.read_csv(get_param_in_cwd('label_file'), dtype={'ID': str})\n",
    "for ug in get_param_in_cwd('subsets'):\n",
    "    sub_group = label_data[label_data['group'] == ug]\n",
    "    display(sub_group)\n",
    "    sub_features = pd.merge(all_2048_features, sub_group[['ID']], on='ID', how='inner')\n",
    "    print(ug, len(np.unique(sub_group['ID'])), len(np.unique(sub_features['ID'])))\n",
    "    sub_features.to_csv(f'{model_root}/{ug}.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7050436a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from onekey_algo import get_param_in_cwd\n",
    "from onekey_algo.fusion.MultiTransformer.run_model import train_categorical_model as clf_main\n",
    "from collections import namedtuple\n",
    "\n",
    "data_root = os.path.join(get_param_in_cwd('model_root'), 'semi_models', get_param_in_cwd('sel_model')['semi'])\n",
    "# 单中心设置参数\n",
    "train = os.path.join(data_root, 'train.csv')\n",
    "val = os.path.join(data_root, 'test.csv')\n",
    "tests = []\n",
    "\n",
    "# 多中心设置参数\n",
    "# train = os.path.join(data_root, 'train.csv')\n",
    "# val = os.path.join(data_root, 'val.csv')\n",
    "# tests = [os.path.join(data_root, subset + '.csv') for subset in get_param_in_cwd('subsets')\n",
    "#          if subset not in ['train', 'val'] and os.path.exists(os.path.join(data_root, subset + '.csv'))]\n",
    "print(tests)\n",
    "# tests = [os.path.join(data_root, 'NHH-GC.csv')]\n",
    "target_file = get_param_in_cwd('label_file')\n",
    "input_dim = features.shape[1] - 1\n",
    "normalize = False\n",
    "header = 0\n",
    "\n",
    "bags_size = 32\n",
    "for bags_size in [bags_size]:\n",
    "    params = dict(train=train,\n",
    "                  val=val,\n",
    "                  tests=tests,\n",
    "                  target_file=target_file,\n",
    "                  j=6,\n",
    "                  input_dim=input_dim,\n",
    "                  bags_size=bags_size,\n",
    "                  normalize=normalize,\n",
    "                  header=header,\n",
    "                  gpus=[0],\n",
    "                  batch_size=16,\n",
    "                  model_name='Transformer',\n",
    "                  epochs=get_param_in_cwd('transformer_epoch', 100),\n",
    "                  init_lr=0.001,\n",
    "                  optimizer='sgd',\n",
    "                  model_root=os.path.join(get_param_in_cwd('model_root'), f'Transformer'),\n",
    "                  add_date=False,\n",
    "                  retrain='',\n",
    "                  iters_start=0,\n",
    "                  iters_verbose=4,\n",
    "                  save_per_epoch=False,\n",
    "                  pretrained=True)\n",
    "    # 训练模型\n",
    "    Args = namedtuple(\"Args\", params)\n",
    "    clf_main(Args(**params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6e520b6",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from onekey_algo.custom.components import metrics\n",
    "from onekey_algo.custom.components.comp1 import draw_roc\n",
    "from onekey_algo import get_param_in_cwd\n",
    "import numpy as np\n",
    "\n",
    "def get_log(epoch):\n",
    "    log_ = pd.concat([pd.read_csv(os.path.join(root, f'../train/Epoch-{epoch}_spec.csv'), dtype={'ID': str}), \n",
    "                      pd.read_csv(os.path.join(root, f'../valid/Epoch-{epoch}_spec.csv'), dtype={'ID': str})], axis=0)\n",
    "    log_.columns = ['ID', 'label-0', 'label-1']\n",
    "    log_['ID'] = log_['ID'].astype(str)\n",
    "    return log_\n",
    "\n",
    "os.makedirs('img', exist_ok=True)\n",
    "label_data = pd.read_csv(get_param_in_cwd('label_file'), dtype={'ID': str})\n",
    "ep_map = { bagsize: get_param_in_cwd('sel_epoch')['transformer']}\n",
    "metrics_df = []\n",
    "dim = 1024\n",
    "for dim in ep_map.keys():\n",
    "    model_root = os.path.join(get_param_in_cwd('model_root'), f'Transformer')\n",
    "    root = os.path.join(model_root, 'Transformer', 'viz')\n",
    "    all_gt = []\n",
    "    all_pred = []\n",
    "    metric_df = []\n",
    "    for idx, subset in enumerate(get_param_in_cwd('subsets')):\n",
    "        sub_group = get_log(ep_map[dim])\n",
    "        sub_group.columns = ['ID', 'label-0', 'label-1']\n",
    "        sub_group = pd.merge(sub_group, label_data[label_data['group'] == subset], on='ID', how='inner')\n",
    "        sub_group[['ID', 'label-0', 'label-1']].to_csv(f'results/Transformer_Transformer_{subset}.csv', index=False)\n",
    "        all_gt.append(np.array(sub_group['label']))\n",
    "        all_pred.append(np.array(sub_group['label-1']))\n",
    "        acc, auc, ci, tpr, tnr, ppv, npv, _, _, _, thres = metrics.analysis_pred_binary(np.array(sub_group['label']), \n",
    "                                                                                        np.array(sub_group['label-1']))\n",
    "        ci = f\"{ci[0]:.4f}-{ci[1]:.4f}\"\n",
    "        metric_df.append([acc, auc, ci, tpr, tnr, ppv, npv, thres, subset])\n",
    "    metric_df = pd.DataFrame(metric_df, columns=['Acc', 'AUC', '95% CI', 'Sensitivity', 'Specificity', 'PPV', 'NPV', 'Youden', 'Cohort'])\n",
    "    metric_df['Dim'] = dim\n",
    "    metrics_df.append(metric_df)\n",
    "    draw_roc(all_gt, all_pred, labels=get_param_in_cwd('subsets'), title=f\"Transformer {dim} Dimension CNN Features\")\n",
    "    plt.savefig(f'img/Transformer_Transformer_roc.svg', bbox_inches='tight')\n",
    "    plt.show()\n",
    "metrics_df = pd.concat(metrics_df, axis=0)\n",
    "metric_df.to_csv('results/Transformer_cmp.csv', index=False)\n",
    "metrics_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "518ed8a3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
